package cn.feizhu.aimarket.ai.workflow.node;

import cn.feizhu.aimarket.ai.client.DocGenerateClientFactory;
import cn.feizhu.aimarket.ai.workflow.model.constant.WorkflowControlConstant;
import cn.feizhu.aimarket.ai.workflow.model.enums.NodeStatus;
import cn.feizhu.aimarket.model.dto.ai.outline.ChapterOutline;
import cn.feizhu.aimarket.model.dto.ai.outline.NovelOutlineResponse;
import cn.hutool.json.JSONUtil;
import com.alibaba.cloud.ai.graph.OverAllState;
import com.alibaba.cloud.ai.graph.action.NodeAction;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.client.ChatClient;

import java.util.*;

/**
 * 多文档生成节点
 * 基于审核通过的大纲生成多个分布式文档
 *
 * @author feizhu
 * @since 2025/10/16
 */
public class MultiDocGenerateNode implements NodeAction {

    private static final Logger logger = LoggerFactory.getLogger(MultiDocGenerateNode.class);

    // TODO 一个客户端只能阻塞等待，这里考虑传入一个工厂
    private final DocGenerateClientFactory docGenerateClientFactory;
    private final Map<String, NodeStatus> node2Status;
    
    public static final String NODE_NAME = "multi_doc_generate";

    private static final String USER_PROMPT = """
            当前章节任务:
            {section_outline}
            
            整体故事背景（仅供参考，不得偏离）:
            {full_outline}
            """;

    public MultiDocGenerateNode(DocGenerateClientFactory docGenerateClientFactory, Map<String, NodeStatus> node2Status) {
        this.docGenerateClientFactory = docGenerateClientFactory;
        this.node2Status = node2Status;
    }

    // TODO 通过多线程改造

    @Override
    public Map<String, Object> apply(OverAllState state) {
        node2Status.put(NODE_NAME, NodeStatus.RUNNING);
        logger.info("=== 多文档生成节点开始执行 ===");

        try {
            // 获取审核通过的大纲 - 处理类型转换问题
            Object outlineObj = state.value(WorkflowControlConstant.OUTLINE).orElse(null);
            if (outlineObj == null) {
                throw new RuntimeException("未找到大纲数据");
            }

            NovelOutlineResponse outline;
            if (outlineObj instanceof NovelOutlineResponse) {
                // 如果已经是正确类型，直接使用
                outline = (NovelOutlineResponse) outlineObj;
            } else if (outlineObj instanceof Map) {
                // 如果是Map类型（LinkedHashMap），转换为JSON再解析为对象
                String jsonStr = JSONUtil.toJsonStr(outlineObj);
                outline = JSONUtil.toBean(jsonStr, NovelOutlineResponse.class);
                logger.info("已将LinkedHashMap转换为NovelOutlineResponse对象");
            } else {
                throw new RuntimeException("不支持的大纲数据类型: " + outlineObj.getClass().getName());
            }

            logger.info("开始基于大纲生成多个文档");

            // 将大纲分割成多个部分进行并行生成
            List<ChapterOutline> chapters = outline.getChapters();
            List<String[]> generatedDocsMap = new ArrayList<>();
            // 用docMap控制有序 因为现在无序
            List<Thread> threads = new ArrayList<>();
            // 为每个部分生成对应的文档内容
            for (int i = 0; i < chapters.size(); i++) {
                // TODO 使用final作为计数，是不是会有问题这里？
                final int currentIndex = i;
                Thread thread = Thread.startVirtualThread(() -> {
                    logger.info("正在生成第{}个文档部分", currentIndex + 1);
                    ChatClient docGenerateClient = docGenerateClientFactory.createDocGenerateClient();
                    String docContent = generateDocumentSection(docGenerateClient, outline, currentIndex + 1);
                    synchronized (generatedDocsMap) {
                        generatedDocsMap.add(new String[]{String.valueOf(currentIndex), docContent});
                    }
                    logger.info("第{}个文档部分生成完成，长度: {}", currentIndex + 1, docContent.length());
                });
                threads.add(thread);
            }
            // 等待所有线程完成
            for (Thread thread : threads) {
                thread.join();
            }
            // 这里对List进行排序
            List<String> generatedDocs = new ArrayList<>();
            generatedDocsMap.stream()
                    .sorted(Comparator.comparingInt(arr -> Integer.parseInt(arr[0])))
                    .forEachOrdered(arr -> generatedDocs.add(arr[1]));

            logger.info("所有文档部分生成完成，共生成{}个文档", generatedDocs.size());
            for (int i = 0; i < generatedDocs.size(); i++) {
                // TODO 先停止预览
//                logger.info("文档部分 {} 内容预览: \n{}", i + 1,
//                        generatedDocs.get(i));
            }
            // 将生成的文档列表存储到状态中
            node2Status.put(NODE_NAME, NodeStatus.COMPLETED);
            return Map.of(WorkflowControlConstant.SUB_DOCS, generatedDocs);

        } catch (Exception e) {
            logger.error("多文档生成节点执行异常", e);
            node2Status.put(NODE_NAME, NodeStatus.FAILED);
            List<String> errorDocs = List.of("文档生成失败：" + e.getMessage());
            return Map.of(WorkflowControlConstant.SUB_DOCS, errorDocs);
        } finally {
            logger.info("=== 多文档生成节点执行完成 ===");
        }
    }


    /**
     * 生成单个文档部分
     */
    private String generateDocumentSection(ChatClient docGenerateClient, NovelOutlineResponse novelOutlineResponse, int sectionNumber) {
        try {
            String content = docGenerateClient.prompt()
                    .user(USER_PROMPT.replace("{section_outline}", JSONUtil.toJsonStr(novelOutlineResponse.getChapters().get(sectionNumber - 1)))
                            .replace("{full_outline}", JSONUtil.toJsonStr(novelOutlineResponse)))
                    .call()
                    .content();

            return String.format("=== 文档部分 %d ===\n%s\n", sectionNumber, content);

        } catch (Exception e) {
            logger.error("生成文档部分{}时发生异常", sectionNumber, e);
            return String.format("=== 文档部分 %d ===\n生成失败：%s\n", sectionNumber, e.getMessage());
        }
    }
}
